fix(lora): export PEFT adapters through Megatron-Bridge - #2580
Conversation
525693b to
32bdd75
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Arist12
changed the title
Keep the stacked checkpoint tree aligned with radixark#2580's focused test registration.
Keep the test-only PR aligned with radixark#2580's simplified coverage and nightly CI scope.
fd3139d to
8be9f3a
Compare
…hards
The staged-error refactor moved the Megatron-native shard collection out of
expose_adapter_slot. Outside that context named_parameters() yields
".adapters.{slot}." names, which the '".adapter." in name' filter never
matches, so every multi-LoRA native shard was written as a valid but empty
file and resume silently restarted from a fresh adapter.
Bridge documents the dependency: native save/load walk named_parameters()
and want names without the slot index so saving from slot A and loading into
slot B produces matching keys.
Treat the Bridge export as an all-or-nothing artifact pair so a failed safetensors write cannot leave a config-only adapter that downstream loaders mistake for valid.
Resolve the writer rank directly so incomplete-export cleanup remains valid after native checkpoint dedup moves shard ownership out of this function.
Write Bridge output in a temporary directory and publish weights before config so failed writes never expose a config-only adapter as loadable.
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
I reviewed the current head (99a218d5028f50a327de0f9923c58fab12242bf1) and tested both successful PEFT loading and rank-local failure behavior on AMD Instinct MI355X GPUs.
The staged safetensors produced by the happy path were loadable through PeftModel.from_pretrained, but the distributed failure protocol still has a merge blocker.
Blocker — staging-directory creation occurs outside setup consensus.
TemporaryDirectory(...) is entered before the inner try that assigns hf_export_err. If it fails on one rank, that rank skips all_gather_object and proceeds toward the trailing barrier while peers wait in all_gather_object. Injecting this failure on one of two MI355 ranks reproducibly timed out after 30 seconds.
Move staging-directory creation into the coordinated setup phase so every rank reaches the same consensus before any rank enters Bridge.
Publication occurs before all-rank save success is known.
Rank 0 moves weights/config into the public directory before save_errors is gathered. A peer failure can therefore be discovered only after files are visible. Overwriting an existing adapter also exposes new weights alongside the old config between the two replacements, and a config-rename failure can make that mismatch permanent.
Use an explicit two-phase protocol:
- create staging and gather setup errors;
- perform the collective Bridge export;
- gather all save results;
- publish on rank 0 only after global success;
- gather promotion results.
A versioned directory or explicit commit marker is safer for overwrite/retry semantics.
The recovery protocol depends on undocumented Bridge internals.
The compensating dist.barrier() assumes save_hf_adapter has exactly one entry and one exit barrier. Its public contract only says the method is collective. The object consensuses also use the default RCCL group, unlike the Gloo helper used elsewhere on this path. Please centralize failure consensus on get_gloo_group() and avoid encoding Bridge's private barrier count.
Finally, keep Bridge's empty-adapter guard on the multi-LoRA helper path; otherwise an empty safetensors file plus target_modules: [] can be promoted as a complete checkpoint.
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
Follow-up review on the unchanged head (99a218d5028f50a327de0f9923c58fab12242bf1): one new retry/topology issue remains after excluding the first review's publication findings.
A failed multi-LoRA attempt directory is reused and later promoted wholesale.
save_multi_lora_checkpoints uses the deterministic _tmp_step_{iteration} directory. Several new coordinated failure points can raise after native shards have already been written, but none removes that directory. On restart, save_due_adapter_checkpoints retries because only final_dir suppresses a save, mkdir(exist_ok=True) reuses the old temporary directory, and os.replace(tmp_dir, final_dir) publishes every stale file that remains in it.
Concrete failure sequence:
- an EP4 attempt writes
_ep0through_ep3shards and then fails during PEFT conversion; - the run restarts at EP1, writes the no-EP-suffix shard into the same temporary directory, and promotes it;
- a later EP4 load sees the stale
_ep0through_ep3set as complete and resumes those pre-failure weights while reporting the retried step as valid.
Use a per-attempt unique staging directory whose name is chosen once and broadcast to all ranks. Alternatively, have one coordinator remove the deterministic temporary directory, reach a barrier, and only then let shard writers recreate/populate it. Concurrent shard writers should not each race to remove the shared directory.
…dapters Create the staging directory inside the guarded block so a rank that fails there cannot skip the consensus and hang its peers, report both consensuses over gloo rather than the GPU communicator that may already be poisoned, restore the pre-promotion os.sync(), and reject an empty multi-LoRA export instead of promoting an empty safetensors with target_modules: [].
|
Thanks. Three points adopted; on the publication-ordering redesign I read Bridge's source and came to a different conclusion — details below so it can be checked. Adopted. The export now stages into a temp directory inside the failure consensus rather than after it, so a rank that fails mid-export cannot leave a half-written directory that a later run mistakes for complete. The consensus all-gather runs on the gloo group rather than the default group, so it does not serialise behind NCCL work. And the empty-adapter case now raises instead of promoting an empty Adopted (comment only). The Not adopted — the 5-step publication protocol / dropping the compensating barrier. The concern is that the extra
Removing the compensating barrier means the non-writer ranks return from Bridge while rank 0 is still writing, then reach the gloo consensus and observe an incomplete directory — which is the failure the barrier exists to prevent. Restructuring the whole publication into a 5-step protocol would mean reimplementing what Bridge already does rather than calling it, which is the opposite of what this PR is for (title: export PEFT adapters through Megatron-Bridge). Also relevant: completion is already signalled by Happy to revisit if I've misread the barrier placement — that is the load-bearing claim. Not adopted — "this PR bundles four changes." The four hunks are the single change of routing PEFT export through Bridge: key prefixing, rank pattern, target modules and dtype are all what Deferred — Validation: |
2 participants
Part of #2705.
Problem
Miles manually assembles PEFT weights and metadata, duplicating Megatron-Bridge conversion logic and producing names or packed-expert layouts that external PEFT loaders may not consume. Multi-adapter writer failures also need rank-aligned handling, and a failed Bridge write must not expose a config-only adapter as loadable.
Change
Delegate single-adapter export to
AutoBridge.save_hf_adapter. Stage Bridge output in a temporary directory, then promote safetensors before config so only a complete pair becomes visible. For Multi-LoRA, slice each slot to its configured rank and use Bridge conversion/configuration helpers before writing safetensors. Coordinate setup, conversion, write, and promotion failures across ranks.Notes for review:
dist.barrier()stays on the default group because it must match Bridge's own exit barrier.target_modules: []as a complete checkpoint.Known boundary: dense and per-expert MoE export work;
--experts-shared-outer-lorasremains native-checkpoint-only because Bridge currently rejects that PEFT layout.Validation
iter_14, exported complete PEFT pairs at 19/24/29, and the final artifact contains 504 finite tensors with no staging residue